# -*- coding: UTF-8 -*-
'''
@Project ：push_rk
@File ：model.py
@IDE ：PyCharm
@Author ：苦瓜
@Date ：2025/9/25 8:34
@Note: Something beautiful is about to happen !
'''
import matplotlib.pyplot as plt
import numpy as np
from keras.src.utils import to_categorical
from keras.src.legacy.preprocessing.image import ImageDataGenerator
from keras import Model, Sequential, optimizers, losses, metrics
from keras.src.layers import Flatten, Dense, Dropout, Conv2D, MaxPooling2D
from keras.src.datasets import mnist

class VGG16(Model):

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        # 卷积层
        self.conv = Sequential([
            Conv2D(filters=16, kernel_size=3, padding='same', activation='relu'),
            Conv2D(filters=16, kernel_size=3, padding='same', activation='relu'),
            MaxPooling2D(),
            Conv2D(filters=32, kernel_size=3, padding='same', activation='relu'),
            Conv2D(filters=32, kernel_size=3, padding='same', activation='relu'),
            MaxPooling2D(),
            Conv2D(filters=64, kernel_size=3, padding='same', activation='relu'),
            Conv2D(filters=64, kernel_size=3, padding='same', activation='relu'),
            Conv2D(filters=64, kernel_size=3, padding='same', activation='relu'),
            MaxPooling2D(),
            Conv2D(filters=128, kernel_size=3, padding='same', activation='relu'),
            Conv2D(filters=128, kernel_size=3, padding='same', activation='relu'),
            Conv2D(filters=128, kernel_size=3, padding='same', activation='relu'),
            MaxPooling2D()
        ])
        # 展平层
        self.flatten = Flatten()
        # 全连接层
        self.fc = Sequential([
            Dense(units=64, activation='relu'),
            Dropout(0.3),
            Dense(units=64, activation='relu'),
            Dropout(0.3),
            Dense(units=10, activation='softmax')
        ])

    def call(self, inputs, *args, **kwargs):
        x = self.conv(inputs)
        x = self.flatten(x)
        x = self.fc(x)
        return x

class MyBaseCNN:

    def __init__(self):
        self.model = None
        self.history = None
        self.y_dim = None

    def load_data(self, datasets_module):
        """ 加载数据集 """
        (X_train, y_train), (X_test, y_test) = datasets_module.load_data()  # 加载数据集
        # 重塑为期望形状
        X_train, X_test = X_train.reshape(-1, 28, 28, 1), X_test.reshape(-1, 28, 28, 1)
        # 进行手动的归一化
        X_train, X_test = X_train.astype('float32') / 255., X_test.astype('float32') / 255.
        # 获取标签类别数量
        y_dim = len(np.unique(y_train))
        # 对标签进行独热
        y_train, y_test = self.one_hot_encoding(y_train, y_dim), self.one_hot_encoding(y_test, y_dim)
        return X_train, X_test, y_train, y_test

    def one_hot_encoding(self, x, cls_num=None):
        """ 对目标数据进行独热 """
        return to_categorical(x, cls_num)

    def dataImagePlus(self, train_dir):
        datagen = ImageDataGenerator(
            rotation_range=25,
            shear_range=0.2,
            width_shift_range=0.1,
            height_shift_range=0.1,
            horizontal_flip=True,
            vertical_flip=True
        )
        return datagen.flow(*train_dir)

    def model_build_compile(self, X_train, y_train, boundary, batch_size=100, epochs=10, validation_data=None):
        """ 模型的构建、编译、训练 """
        self.model = VGG16()
        self.model.build(input_shape=(None, boundary, boundary, 1))  # 模型构建
        # 模型编译
        self.model.compile(
            optimizer=optimizers.Adam(),
            loss=losses.categorical_crossentropy,
            metrics=['acc', metrics.Recall, metrics.Precision, metrics.F1Score]
        )
        # 模型训练
        self.history = self.model.fit(
            self.dataImagePlus(train_dir=(X_train, y_train)),
            batch_size=batch_size, epochs=epochs,
            validation_data=validation_data
        )
        return self.model, self.history

def main():
    # 实例基础模型
    base_obj = MyBaseCNN()
    # 获取数据
    X_train, X_test, y_train, y_test = base_obj.load_data(mnist)
    # 构建模型
    model, history = base_obj.model_build_compile(X_train, y_train, boundary=X_train.shape[1],
                                                  validation_data=(X_test,y_test))
    # 输出模型形状
    model.summary()
    # 评估模型
    eval_dict = model.evaluate(X_test, y_test, return_dict=True)
    print(eval_dict)
    # 可视化模型训练和验证的准确率曲线
    plt.plot(history.history['acc'], label='train acc', c='r')
    plt.plot(history.history['val_acc'], label='val acc', c='g')
    plt.legend()
    plt.show()
    # 保存模型权重
    model.save_weights('model.weights.h5')

if __name__ == '__main__':
    main()

# 1)	数据预处理约束：
# MNIST手写数字数据集被划分为训练集和测试集。
# 输入图像数据（x_train和x_test）被重新塑形为四维张量，并归一化到0到1之间。
# 图像数据增强技术的应用，实现随机翻转、旋转、裁剪等图像数据增强。
# 2)	模型架构约束：
# 采用模型进行构建，包含卷积层、池化层、扁平化层和全连接层。
# 卷积层使用0填充，以保持特征图的尺寸不变。
# 池化层用于降低特征图的维度。
# 扁平化层将多维特征图转换为一维向量，以便输入到全连接层。
# 全连接层包括两个具有激活函数的隐藏层，以及一个具有激活函数的输出层，用于多分类任务。
# Dropout层被用于减少过拟合，设置丢弃率。
# 3)	模型编译约束：
# 使用优化器进行模型训练。
# 损失函数选择为交叉熵损失函数，适用于多分类问题且标签为整数形式。
# 监控模型的准确率作为性能指标。
# 4)	模型训练约束：
# 模型通过fit方法进行训练，指定了训练轮数和批量大小。
# 使用了验证集来评估模型在未见数据上的性能。
# 5)	训练过程可视化约束：
# 绘制训练和验证的准确率曲线，以便直观观察模型性能。
# 6)	模型保存约束：
# 训练完成后，保存模型的权重到指定文件。
